import cv2
import numpy as np
import time
from gif import path_point_save


def marching_squares(image):
    """
    marching squares 算法

    物体必须位于图像中，不能和图像边缘相贴合

    :param image:
    :return:
    """
    h, w = image.shape

    points_list = []
    start_point = (1, 1)
    is_visit = np.zeros_like(image, dtype=np.bool)  # 用于图像存在多物体时
    while 1:
        start_point = get_start_point(image, is_visit, start_point)
        if start_point is None:
            break

        points = [start_point]
        is_visit[start_point[1], start_point[0]] = True
        p_x, p_y = start_point
        step_x, step_y = None, None
        prev_x, prev_y = None, None
        while 1:
            sv = get_square_value(image, (p_x, p_y))

            if sv == 1 or sv == 5 or sv == 13:
                '''
                going UP with these cases:
                +---+---+   +---+---+  +---+---+
                | 1 |   |   | 1 |   |   | 1 |   |
                +---+---+   +---+---+  +---+---+
                |   |   |   | 4 |   |    | 4 | 8 |
                +---+---+  +---+---+   +---+---+
                '''
                step_x, step_y = 0, -1

            elif sv == 8 or sv == 10 or sv == 11:
                '''
                going DOWN with these cases:
                +---+---+   +---+---+  +---+---+
                |   |   |   |   | 2 |   | 1 | 2 |
                +---+---+  +---+---+  +---+---+
                |   | 8 |   |   | 8 |   |   | 8 |
                +---+---+  +---+---+   +---+---+
                '''
                step_x, step_y = 0, 1

            elif sv == 4 or sv == 12 or sv == 14:
                '''
                going LEFT with these cases:
                +---+---+   +---+---+   +---+---+
                |   |   |    |   |   |     |   | 2 |
                +---+---+   +---+---+   +---+---+
                | 4 |   |   | 4 | 8 |   | 4 | 8 |
                +---+---+   +---+---+   +---+---+
                '''
                step_x, step_y = -1, 0

            elif sv == 2 or sv == 3 or sv == 7:
                '''
                going RIGHT with these cases:
                +---+---+   +---+---+   +---+---+
                |   | 2 |   | 1 | 2 |   | 1 | 2 |
                +---+---+   +---+---+   +---+---+
                |   |   |    |   |   |    | 4 |   |
                +---+---+   +---+---+   +---+---+
                '''
                step_x, step_y = 1, 0

            elif sv == 6:
                '''
                special saddle point case 1:

                +---+---+
                |   | 2 |
                +---+---+
                | 4 |   |
                +---+---+

                going LEFT if coming from UP
                else going RIGHT
                '''
                if prev_x == 0 and prev_y == -1:
                    step_x, step_y = -1, 0
                else:
                    step_x, step_y = 1, 0

            elif sv == 9:
                '''
                special saddle point case 2:

                +---+---+
                | 1 |   |
                +---+---+
                |   | 8 |
                +---+---+

                going UP if coming from RIGHT
                else going DOWN
                '''
                if prev_x == 1 and prev_y == 0:
                    step_x, step_y = 0, -1
                else:
                    step_x, step_y = 0, 1

            p_x, p_y = p_x + step_x, p_y + step_y
            prev_x, prev_y = step_x, step_y

            # 回到起始点
            if p_x == start_point[0] and p_y == start_point[1]:
                break
            # 边界限制（超出图像
            if p_x == w or p_y == h:
                break
            points.append((p_x, p_y))
            is_visit[p_y, p_x] = True

            # 边界限制（下一步不能4合1
            if p_x == 0 or p_y == 0:
                break

        points_list.append(points)
    return points_list


def get_start_point(img, is_visit, _point:tuple=(1,1)):
    h, w = img.shape
    start_x, start_y = _point
    for i in range(start_y, h):
        for j in range(0, w):
            if is_visit[i, j]:
                continue
            _sum = np.sum(img[i-1:i+1, j-1:j+1])
            if 0 < _sum < 4:
                return (j, i)
    return None


def get_square_value(img, _point: tuple|list):
    """
        +-----+----+
        |  1  |  2  |
        +-----+----+
        |  4  |  8  |  <--- current point (pX, pY)
        +-----+----+
    """
    val = 0
    _x, _y = _point
    if img[_y-1, _x-1] == 0:
        val = 1
    if img[_y-1, _x] == 0:
        val += 2
    if img[_y, _x-1] == 0:
        val += 4
    if img[_y, _x] == 0:
        val += 8
    return val


def marching_squares_exp(image):
    """
    marching squares 算法扩展

    Parameters
    ----------
    image : np.ndarray
        二值图，0 or 1，0是物体，1是背景
    """
    # 最外围填充一圈1。确保物体处于图像中，而不是被边缘裁切
    pad_size = 1
    pad_img = np.pad(image, pad_size, 'constant', constant_values=1)

    # 使用marching_squares算法获取轮廓路线
    paths = marching_squares(pad_img)

    h, w = image.shape
    fact_paths = []
    for points in paths:
        points = np.array(points)
        new_paths = [[]]

        # 矫正偏移
        points = points - pad_size

        # 移除不属于原始图像的路径点
        for i in range(len(points)):
            x, y = points[i]
            is_outside = ((x <= 0) | (x >= w) | (y <= 0) | (y >= h)).item()
            if is_outside is False:
                new_paths[-1].append(points[i])
            elif len(new_paths[-1]) > 0:
                new_paths.append([])
        if len(new_paths[-1]) == 0:
            new_paths.pop()

        # 尝试合并首、尾路线
        if len(new_paths) > 1:
            head_start_point = new_paths[0][0]
            tail_end_point = new_paths[-1][-1]
            l2_dist = np.linalg.norm(head_start_point - tail_end_point)
            if l2_dist < 2:
                new_paths[-1].extend(new_paths[0])
                new_paths = new_paths[1:]

        # added final path list
        fact_paths.extend(new_paths)
    return fact_paths


def get_contour_eqdist_path(img):
    """
    生成轮廓等距路径

    Parameters
    -----------
    img : np.ndarray
        二值图，0和255，或者0和1

    """
    dst = np.empty(img.shape, dtype=np.float32)
    cv2.distanceTransform(img, cv2.DIST_L2, 3, dst)

    # 归一化
    temp = cv2.normalize(dst, None, 0, 255, cv2.NORM_MINMAX, cv2.CV_8U)
    cv2.imwrite('output/distanceTransform_normalize.png', temp)
    dict_img = cv2.normalize(dst, None, 0, np.ceil(dst.max()), cv2.NORM_MINMAX, cv2.CV_8U)

    # 使用marching square算法获取轮廓
    final_path_list = []
    for level in range(2, dict_img.max(), 3):
        mask = dict_img < level
        paths = marching_squares_exp(mask)
        final_path_list.extend(paths)

    path_point_save(img, final_path_list, "./output", "contour_test")
    return final_path_list


def load_image(image_path: str):
    # 以灰度模式读取图片
    image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    # 应用阈值进行二值化
    _, binary_image = cv2.threshold(image, 150, 255, cv2.THRESH_BINARY)

    return binary_image


def main():
    # 图片路径
    image_path = "images/marching_squares_004.png"
    # 加载二值图
    binary_image = load_image(image_path)
    #
    get_contour_eqdist_path(binary_image)

def ms_test(file_name):
    image = load_image(f"images/{file_name}.png")
    image[image == 255] = 1
    paths = marching_squares_exp(image)
    path_point_save(image, paths, "./output", file_name)


if __name__ == '__main__':
    start_time = time.perf_counter()

    main()

    # file_name = "marching_squares_004"
    # ms_test(file_name)


    end_time = time.perf_counter()
    elapsed_time_ms = (end_time - start_time) * 1000
    print(f'耗时: {elapsed_time_ms:.3f}ms, done.')
